#include "poseestimation.h"
#include "pfmimage.h"
#include <QFile>
#include <QTextStream>
#include <QMessageBox>
#include <QGuiApplication>
#include <fstream>
#include <string>
PoseEstimation::PoseEstimation()
{
    imageNum = 0;
    fiducialNum = 9;
    oriImgNum = 0;
    meshFiducialsList.resize(fiducialNum);
    logTextStream = NULL;
    iterNum = 0;
}
PoseEstimation::~PoseEstimation()
{

}

void PoseEstimation::SetLogTextStream(QTextStream *p)
{
    logTextStream = p;
}

void PoseEstimation::OutputMatrixToLog(cv::Mat cvMat, QString matName)
{
    if(logTextStream)
    {
        *logTextStream<<matName<<"=[\n";
        short rows = cvMat.rows;
        short cols = cvMat.cols;
        for(short i=0; i<rows; i++)
        {
            for(short j=0; j<cols; j++)
            {
                *logTextStream<<cvMat.at<float>(i,j)<<"\t";
            }
            *logTextStream<<"\n";
        }
        *logTextStream<<"];\n";
        logTextStream->flush();
    }
}

void PoseEstimation::ComputeCentroidImageFiducials()
{
    centroidImageFiducialsList.clear();
    centroidImageFiducialsList.resize(imageNum);
    float centroidX;
    float centroidY;
    for(int i=0; i<imageNum; i++)
    {
        centroidX = 0.0f;
        centroidY = 0.0f;

        for(int j=0; j<fiducialNum; j++)
        {
            centroidX += imageFiducialsList[i*fiducialNum+j].x();
            centroidY += imageFiducialsList[i*fiducialNum+j].y();
        }
        centroidImageFiducialsList[i].setX(centroidX/fiducialNum);
        centroidImageFiducialsList[i].setY(centroidY/fiducialNum);
    }
}

void PoseEstimation::ComputeCentroidMeshFiducials()
{

    float centroidX, centroidY, centroidZ;
    centroidX = 0.0f;
    centroidY = 0.0f;
    centroidZ = 0.0f;
    for(int i=0; i<fiducialNum; i++)
    {
        centroidX += meshFiducialsList[i].x();
        centroidY += meshFiducialsList[i].y();
        centroidZ += meshFiducialsList[i].z();
    }
    centroidX /= fiducialNum;
    centroidY /= fiducialNum;
    centroidZ /= fiducialNum;

    centroidMeshFicucial.setX(centroidX);
    centroidMeshFicucial.setY(centroidY);
    centroidMeshFicucial.setZ(centroidZ);

}


bool PoseEstimation::ReadImageFiducials(QString fileName)

{


    imageFilePath = fileName;
    short tmpIndex = imageFilePath.lastIndexOf("/");
    imageFilePath.remove(tmpIndex+1,imageFilePath.size()-tmpIndex-1);
    std::ifstream infile;
    infile.open(fileName.toStdString());
    if(!infile.is_open())
        return false;

    infile>>oriImgNum;

    std::string imageName;
    float tmpValue;
    short flag;
    usedImageList.clear();

    for(int i=0; i<oriImgNum; i++)
    {
        infile>>flag;
        if(flag == -1)
        {
            infile>>imageName;
            for(int k=0; k<22; k++)
                infile>>tmpValue;
            continue;
        }

        infile>>imageName;
        usedImgFileNameList.append(QString::fromStdString(imageName));

        QVector2D tmpfaceCenter;
        infile>>tmpValue;
        tmpfaceCenter.setX(tmpValue);
        infile>>tmpValue;

        tmpfaceCenter.setY(tmpValue);
        faceCenterList.append(tmpfaceCenter);

        infile>>tmpValue;
        faceBoxSizeList.append(tmpValue);
        infile>>tmpValue;
        confidenceList.append(tmpValue);

        for(int j=0; j<fiducialNum; j++)//9 fiducials for each face
        {
            QVector2D tmpFiducial;
            infile>>tmpValue;
            tmpFiducial.setX(tmpValue);
            infile>>tmpValue;

            tmpFiducial.setY(tmpValue);

            imageFiducialsList.append(tmpFiducial);
        }
        usedImageList.append(i+1);

    }
    imageNum = usedImageList.size();

    infile.close();

    this->ComputeCentroidImageFiducials();

    return true;

}

bool PoseEstimation::ReadMeshFiducials(QString fileName, QString meshFileName)
{
    std::ifstream inFile;
    inFile.open(fileName.toStdString());
    if(!inFile.is_open())
        return false;

    PFMImage pfmImage;
    pfmImage.ReadFromFile(meshFileName);

    float u, v;
    for(int i=0; i<fiducialNum; i++)
    {
        inFile>>u>>v;
        meshFiducialsList[i] = pfmImage.GetPixelValue(u, v);
    }
    inFile.close();

    this->ComputeCentroidMeshFiducials();

    return true;
}

void PoseEstimation::ComputePosePara()
{


    cvRotation.clear();
    cvTrans.clear();
    scales.clear();
    cvRotation.resize(imageNum);
    cvTrans.resize(imageNum);
    scales.resize(imageNum);


    //initialize matrix qMat with mesh fiducials
    cv::Mat qMat(3, fiducialNum, CV_32F);
    cv::Mat qTMat(fiducialNum, 3, CV_32F);
    for(int i=0; i<fiducialNum; i++)
    {
       float tx = meshFiducialsList[i].x() - centroidMeshFicucial.x();
       float ty = meshFiducialsList[i].y() - centroidMeshFicucial.y();
       float tz = meshFiducialsList[i].z() - centroidMeshFicucial.z();
       qMat.at<float>(0,i) = tx;
       qMat.at<float>(1,i) = ty;
       qMat.at<float>(2,i) = tz;

       qTMat.at<float>(i,0) = tx;
       qTMat.at<float>(i,1) = ty;
       qTMat.at<float>(i,2) = tz;
    }

//    this->OutputMatrixToLog(qMat, "q");
//    this->OutputMatrixToLog(qTMat, "qt");

    for(int i=0; i<imageNum; i++)
    {

        //initialize matrix pMat with image fiducials

//        if(logTextStream)
//        {
//            *logTextStream<<"image"<<i<<"*******************************************************************************************************\n";
//        }

        cv::Mat pMat(2, fiducialNum, CV_32F);

        for(int k=0; k<pMat.cols; k++)
        {
            float tu = imageFiducialsList[i*fiducialNum + k].x() - centroidImageFiducialsList[i].x();
            float tv = imageFiducialsList[i*fiducialNum + k].y() - centroidImageFiducialsList[i].y();
            pMat.at<float>(0, k) = tu;
            pMat.at<float>(1, k) = tv;
        }

//        this->OutputMatrixToLog(pMat, "p");

        cv::Mat aMat(2,3,CV_32F);
        cv::Mat tmpM = qMat*qTMat;
//        this->OutputMatrixToLog(tmpM, "qqt");
        aMat = pMat*qTMat*tmpM.inv();
//        this->OutputMatrixToLog(aMat, "a");


        cv::Mat pCentroid(2,1,CV_32F);
        pCentroid.at<float>(0,0) = centroidImageFiducialsList[i].x();
        pCentroid.at<float>(1,0) = centroidImageFiducialsList[i].y();
        cv::Mat qCentroid(3,1,CV_32F);
        qCentroid.at<float>(0,0) = centroidMeshFicucial.x();
        qCentroid.at<float>(1,0) = centroidMeshFicucial.y();
        qCentroid.at<float>(2,0) = centroidMeshFicucial.z();
        cv::Mat trans(2,1,CV_32F);
        trans = pCentroid - aMat*qCentroid;
//        this->OutputMatrixToLog(trans, "t");

        cv::Mat row1(1,3, CV_32F);
        cv::Mat row2(1,3, CV_32F);
        row1 = aMat.row(0);//.copyTo(row1.row(1));
        row2 = aMat.row(1);//.copyTo(row2.row(1));
        cv::Mat cross = row1.cross(row2);
        aMat.resize(3);
        row1.row(0).copyTo(aMat.row(0));
        row2.row(0).copyTo(aMat.row(1));
        cross.row(0).copyTo(aMat.row(2));
//        this->OutputMatrixToLog(aMat, "an");

        cv::Mat uMat(3,3,CV_32F);
        cv::Mat dMat(3,1,CV_32F);
        cv::Mat vTMat(3,3,CV_32F);
        cv::SVD::compute(aMat,dMat, uMat, vTMat);
//        this->OutputMatrixToLog(uMat,"u");
//        this->OutputMatrixToLog(vTMat, "vt");
//        this->OutputMatrixToLog(dMat, "d");
        cv::Mat rMat = uMat*vTMat;
//        this->OutputMatrixToLog(rMat, "r");
        //选择两个相近的值取平均作为scale
        float d1=dMat.at<float>(0,0);
        float d2=dMat.at<float>(1,0);
        float d3=dMat.at<float>(2,0);
        float diff12 = abs(d1-d2);
        float diff13 = abs(d1-d3);
        float diff23 = abs(d2-d3);
        float s = 0.0f;
        if(diff12 < diff13)
        {
            if(diff23 < diff12)
            {
                s=(d2+d3)/2.0f;
            }
            else
            {
                s = (d1+d2)/2.0f;
            }
        }
        else
        {
            if(diff23 < diff13)
            {
                s=(d2+d3)/2.0f;
            }
            else
            {
                s = (d1+d3)/2.0f;
            }
        }


        cvRotation[i] = rMat;
        cvTrans[i] = trans;
        scales[i] = s;
    }

    //compute yaw, pitch and roll,
    QVector3D angles;
    float yaw, pitch, roll;
    //angles.x() yaw
    //angles.y() pitch
    //angles.z() roll
    rotAngle.clear();
    for(int i=0; i<imageNum; i++)
    {
        cv::Mat rotMat = cvRotation[i];
        float r00 = rotMat.at<float>(0,0);
        float r10 = rotMat.at<float>(1,0);
        yaw = atan2(r10, r00);
        float r20 = rotMat.at<float>(2,0);
        float r21 = rotMat.at<float>(2,1);
        float r22 = rotMat.at<float>(2,2);
        pitch = atan2(-r20, sqrt(r21*r21+r22*r22));
        roll = atan2(r21, r22);

        yaw = yaw*180/3.1415926;
        pitch = pitch*180/3.1415926;
        roll = roll*180/3.1415926;
        angles.setX(yaw);
        angles.setY(pitch);
        angles.setZ(roll);
        rotAngle.append(angles);
    }
    if(iterNum == 0)
        this->ComputeNearFrontalImgList();

}

void PoseEstimation::ComputeNearFrontalImgList()
{
    float yaw = 0.0f;
    float pitch = 0.0f;

    frontalImgListInUsedImgList.clear();

    for(int i=0; i<imageNum; i++)
    {
        yaw = rotAngle[i].x();
        pitch = rotAngle[i].y();

        if(abs(yaw-180)<5 || abs(pitch)<5)
        {
            frontalImgListInUsedImgList.append(i);
        }
    }
}

void PoseEstimation::ConstructPixelMatrix(QVector<short> imgList, QString geoImgFileName)
{
    if(imageNum==0 || !imgList.size())
        return ;
    int imgListSize = imgList.size();
    QVector<QVector3D> geoPointsList;

    PFMImage geoImg;
    geoImg.ReadFromFile(geoImgFileName);

    short w = geoImg.GetWidth();
    short h = geoImg.GetHeight();

    QVector<QVector2D> pixelPosList;
    for(short u=0; u<w; u++)
    {
        for(short v=0; v<h; v++)
        {
            QVector3D vert = geoImg.GetPixelValue(u,v);
            if(abs(vert.x()) + abs(vert.y()) + abs(vert.z())>0.0)
            {
                geoPointsList.append(vert);
                QVector2D uv;
                uv.setX(u);
                uv.setY(v);
                pixelPosList.append(uv);
            }


        }
    }
    int pointsSize = geoPointsList.size();

    if(iterNum == 0)
        pixelMatrix = cv::Mat::zeros(imgListSize, pointsSize, CV_32F);


    for(int i=0; i<imgListSize; i++)
    {

        QString imgName = imageFilePath + usedImgFileNameList[imgList[i]];
        QImage image(imgName);

        cv::Mat rMat = cvRotation[imgList[i]];
        cv::Mat trans = cvTrans[imgList[i]];
        float scale = scales[imgList[i]];

        for(int j=0; j<pointsSize; j++)
        {
           cv::Mat vertMat(3,1, CV_32F);
           QVector3D vert  = geoPointsList[j];
           vertMat.at<float>(0,0) = vert.x();
           vertMat.at<float>(1,0) = vert.y();
           vertMat.at<float>(2,0) = vert.z();
           cv::Mat ProjVert = scale*rMat*vertMat;

           float u = ProjVert.at<float>(0,0);
           float v = ProjVert.at<float>(1,0);
           u = u + trans.at<float>(0,0);
           v = v + trans.at<float>(1,0);

           //bilinear interpolation
           float w1 = u-(short)u;
           float w2 = v-(short)v;
           short lu = (short)u;
           short lv = (short)v;

           QRgb color1 = image.pixel(lu,lv);
           QColor qColor1 = QColor::fromRgb(color1);
           QRgb color2 = image.pixel(lu+1,lv);
           QColor qColor2 = QColor::fromRgb(color2);
           QRgb color3 = image.pixel(lu,lv+1);
           QColor qColor3 = QColor::fromRgb(color3);
           QRgb color4 = image.pixel(lu+1,lv+1);
           QColor qColor4 = QColor::fromRgb(color4);

           float r1 = qColor1.redF();
           float r2 = qColor2.redF();
           float r3 = qColor3.redF();
           float r4 = qColor4.redF();

           float g1 = qColor1.greenF();
           float g2 = qColor2.greenF();
           float g3 = qColor3.greenF();
           float g4 = qColor4.greenF();

           float b1 = qColor1.blueF();
           float b2 = qColor2.blueF();
           float b3 = qColor3.blueF();
           float b4 = qColor4.blueF();

/*           float gamma = 1/2.2f;
           r1 = pow(r1,gamma);
           r2 = pow(r2,gamma);
           r3 = pow(r3,gamma);
           r4 = pow(r4,gamma);

           g1 = pow(g1,gamma);
           g2 = pow(g2,gamma);
           g3 = pow(g3,gamma);
           g4 = pow(g4,gamma);

           b1 = pow(b1,gamma);
           b2 = pow(b2,gamma);
           b3 = pow(b3,gamma);
           b4 = pow(b4,gamma);
*/
           float r = (r1*(1.0-w1) + r2*w1)*(1.0-w2)+
                     (r3*(1.0-w1) + r4*w1)*w2;
           float g = (g1*(1.0-w1) + g2*w1)*(1.0-w2)+
                     (g3*(1.0-w1) + g4*w1)*w2;
           float b = (b1*(1.0-w1) + b2*w1)*(1.0-w2)+
                     (b3*(1.0-w1) + b4*w1)*w2;

           float grey = 0.299*r + 0.587*g + 0.114*b;

           pixelMatrix.at<float>(i,j) = grey;

        }

        //debug begin
        QImage tmpImage(256, 256, QImage::Format_RGB888);
        QRgb color;

        for(int k=0; k<pixelPosList.size(); k++)
        {
            float value = pixelMatrix.at<float>(i,k);
            int rgb = value*255 + 0.5f;
            color = qRgb(rgb,rgb,rgb);

            int u = pixelPosList[k].x();
            int v = pixelPosList[k].y();
            tmpImage.setPixel(u, v, color);
        }
        imgName.remove(imgName.size()-4, 4);
        imgName += ".png";
        tmpImage.save(imgName, "PNG");
        //debug end
    }

//    this->OutputMatrixToLog(pixelMatrix, "sc");

}

void PoseEstimation::MappingImagetoGeo(short imgIndex, QString geoImgFileName)
{
    if(imageNum==0 || imgIndex > imageNum)
        return ;


    cv::Mat rMat = cvRotation[imgIndex];
    cv::Mat trans = cvTrans[imgIndex];
    float scale = scales[imgIndex];


    QString imgName = imageFilePath + usedImgFileNameList[imgIndex];

    QString outputFileName = usedImgFileNameList[imgIndex];
    short dotIndex = outputFileName.lastIndexOf(".");
    outputFileName.remove(dotIndex, outputFileName.size() - dotIndex);
    outputFileName = imageFilePath + outputFileName + ".pfm";


    QImage image(imgName);
    PFMImage geoImg;
    geoImg.ReadFromFile(geoImgFileName);


    short w = geoImg.GetWidth();
    short h = geoImg.GetHeight();
    PFMImage mapColorImg(w,h,3,0);
    cv::Mat vertMat(3,1, CV_32F);
    QVector3D color;
    for(short u=0; u<w; u++)
    {
        for(short v=0; v<h; v++)
        {
            QVector3D vert = geoImg.GetPixelValue(u,v);
            if(abs(vert.x()) + abs(vert.y()) + abs(vert.z())>0.0)
            {
                vertMat.at<float>(0,0) = vert.x();
                vertMat.at<float>(1,0) = vert.y();
                vertMat.at<float>(2,0) = vert.z();

                cv::Mat ProjVert = scale*rMat*vertMat;

                float u = ProjVert.at<float>(0,0);
                float v = ProjVert.at<float>(1,0);
                u = u + trans.at<float>(0,0);
                v = v + trans.at<float>(1,0);

                //bilinear interpolation
                float w1 = u-(short)u;
                float w2 = v-(short)v;
                short lu = (short)u;
                short lv = (short)v;

                QRgb color1 = image.pixel(lu,lv);
                QColor qColor1 = QColor::fromRgb(color1);
                QRgb color2 = image.pixel(lu+1,lv);
                QColor qColor2 = QColor::fromRgb(color2);
                QRgb color3 = image.pixel(lu,lv+1);
                QColor qColor3 = QColor::fromRgb(color3);
                QRgb color4 = image.pixel(lu+1,lv+1);
                QColor qColor4 = QColor::fromRgb(color4);

                float r = (qColor1.redF()*(1.0-w1) + qColor2.redF()*w1)*(1.0-w2)+
                          (qColor3.redF()*(1.0-w1) + qColor4.redF()*w1)*w2;
                float g = (qColor1.greenF()*(1.0-w1) + qColor2.greenF()*w1)*(1.0-w2)+
                          (qColor3.greenF()*(1.0-w1) + qColor4.greenF()*w1)*w2;
                float b = (qColor1.blueF()*(1.0-w1) + qColor2.blueF()*w1)*(1.0-w2)+
                          (qColor3.blueF()*(1.0-w1) + qColor4.blueF()*w1)*w2;

                color.setX(r);
                color.setY(g);
                color.setZ(b);

            }
            else
            {
                color.setX(0);
                color.setY(0);
                color.setZ(0);
            }
            mapColorImg.SetPixelValue(u,v,color);
        }
    }

    mapColorImg.WriteToFile(outputFileName);
}

QString PoseEstimation::UpdateMeshColor(Mesh* pMesh, short imgIndex)
{
    QString imgFileName;
    if(imageNum==0 || imgIndex > imageNum)
        return imgFileName;

    cv::Mat rMat = cvRotation[imgIndex];
    cv::Mat trans = cvTrans[imgIndex];
    float scale = scales[imgIndex];
    QVector<QVector3D>& verticesList = pMesh->VerticesList();
    float* verticesColorList = pMesh->VerticesColorPointer();

    QString imgName = imageFilePath + usedImgFileNameList[imgIndex];

    imgFileName = usedImgFileNameList[imgIndex];
    short dotIndex = imgFileName.lastIndexOf(".");
    imgFileName.remove(dotIndex, imgFileName.size() - dotIndex);


    QImage image(imgName);

    int verticeNum = verticesList.size();
    cv::Mat vertMat(3,1, CV_32F);
    for(int i=0; i<verticeNum; i++)
    {
        vertMat.at<float>(0,0) = verticesList[i].x();
        vertMat.at<float>(1,0) = verticesList[i].y();
        vertMat.at<float>(2,0) = verticesList[i].z();

        cv::Mat ProjVert = scale*rMat*vertMat;

        float u = ProjVert.at<float>(0,0);
        float v = ProjVert.at<float>(1,0);
        u = u + trans.at<float>(0,0);
        v = v + trans.at<float>(1,0);

        QRgb color = image.pixel((int)(u+0.5f),(int)(v+0.5f));
        QColor qColor = QColor::fromRgb(color);
        float r = qColor.redF();
        float g = qColor.greenF();
        float b = qColor.blueF();

        verticesColorList[i*3] = r;
        verticesColorList[i*3 + 1] = g;
        verticesColorList[i*3 + 2] = b;

    }

    return imgFileName;

}

QVector<short> PoseEstimation::GetFrontalImgIndexList()
{
    QVector<short> list;
    for(int i=0; i<frontalImgListInUsedImgList.size(); i++)
    {
        short index = frontalImgListInUsedImgList[i];
        list.append(usedImageList[index]);
    }
    return list;

}

QVector<QString> PoseEstimation::GetFrontImgFileNameList()
{
    QVector<QString> list;
    for(int i=0; i<frontalImgListInUsedImgList.size(); i++)
    {
        short index = frontalImgListInUsedImgList[i];
        list.append(usedImgFileNameList[index]);
    }
    return list;
}

